import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches


class Orthogonal_Polynomials_Fitting:
    def __init__(self, x, y):
        assert x.shape[0] == y.shape[0] or len(x) == len(y)
        self.x = x
        self.y = y
        self.N = self.x.shape[0]
        self.LP_records = {}
        self.CH_records = {}
        self.low = np.min(x)
        self.high = np.max(x)
        self.shift = False
        if self.low < -1 or self.high > 1:
            self.shift = True
            self.shift_x = np.arange(-1, 1, 2./self.N)

    @classmethod
    def Legender_Polynomials(self, x, n):
        if n == 0:
            return 1
        elif n == 1:
            return x
        else:
            if (x, n) in self.LP_records.keys():
                return self.LP_records[(x, n)]
            else:
                ans = ((2*n-1)*x*Legender_Polynomials(x, n-1) -
                       (n-1)*Legender_Polynomials(x, n-2))/(n)
                self.LP_records[(x, n)] = ans
                return ans

    @classmethod
    def Chebyshev_Polynomials(self, x, n):
        if n == 0:
            return 1
        else:
            return np.cos(n*x)
#         elif n==1:
#             return x
#         else:
#             if (x,n) in self.CH_records.keys():
#                 return self.CH_records[(x,n)]
#             else:
#                 ans=2*x*Chebyshev_Polynomials(x,n-1)-Chebyshev_Polynomials(x,n-2)
#                 self.CH_records[(x,n)]=ans
#                 return ans

    def calculate(self, order, method='Legender', show_error=True, show_fitting_result=True):

        cx = self.shift_x if self.shift else self.x

        if method == 'Legender':
            Method = Legender_Polynomials
        elif method == 'Chebyshev':
            Method = Chebyshev_Polynomials
            # with x=cosθ,we can let the weight function to be 1
            cx = np.arange(0, np.pi, np.pi/cx.shape[0])

        A = np.zeros((order, order))
        for i in range(order):
            for x_ in cx:
                A[i, i] += Method(x_, i)*Method(x_, i)

        B = np.zeros((order, 1))
        for i in range(order):
            for x_, y_ in zip(cx, self.y):
                B[i, 0] += Method(x_, i)*y_

        ai = np.linalg.inv(A)@B
        predict_y = np.zeros_like(self.y)
        for j in range(order):
            for i in range(cx.shape[0]):
                #                 x_ = self.x[i] if not self.shift else 1/(self.high-self.low)*(2*self.x[i]-self.high-self.low)
                predict_y[i] += Method(cx[i], j)*ai[j, 0]

        error = np.linalg.norm(y-predict_y)

        if show_error:
            print("Fitting Error: {:.5f}".format(error))

        if show_fitting_result:
            plt.subplot(111)
            plt.scatter(self.x, self.y, s=5, marker='*', c='green')
            plt.scatter(self.x, predict_y, s=5, marker='o', c='red')
            Raw = mpatches.Patch(color='green', label='Raw Points')
            Predict = mpatches.Patch(color='red', label='Predict Points')
            plt.title(
                'Orthogonal-Polynomials by {} with {} orders'.format(method, order))
            plt.legend(handles=[Raw, Predict], loc='best')
            plt.show()
        return predict_y, error
